import pandas as pd
from openai import OpenAI
from sklearn.metrics import accuracy_score
import ollama
from pandasai.llm.local_llm import LocalLLM
from pandasai import SmartDataframe
import numpy as np
import pandas as pd
from warnings import simplefilter
import streamlit as st
from pandasai.llm import OpenAI as ai
import folium
from folium.plugins import HeatMap

simplefilter(action="ignore", category=pd.errors.PerformanceWarning)
simplefilter(action="ignore", category=FutureWarning)
#ignore ValueError
simplefilter(action="ignore", category=ValueError)
API_key = 'apikey'
client = OpenAI(api_key=API_key)

# question = input('Enter the question you want to ask the model: ')


def openai_result(query, df):
    model = ai(API_key)

    sdf = SmartDataframe(df = df, config={"llm":model})
    try:
        response = sdf.chat(query+'If graphing, you must return a sentence: Graphing completed. Don\'t graph if not specifically asked')
        return response
    except:
        return 'FAILED'
    


def local_llm_result(model,query,df):
    sdf = SmartDataframe(df=df, config={"llm": model,"save_logs": True, "conversational": True,"verbose":True})
    
    result = sdf.chat(query+'If graphing, you must return a sentence: Graphing completed')
    return result

def truth_value(df,query):
    result = eval(query)
    return result


if __name__ == '__main__':
    # initializations
    st.title('OpenAI ChatGPT')

    filename = 'data/cleaned_data.csv'
    df = pd.read_csv(filename, low_memory=False)



    # print(result_llm)
    query_to_openai = 'I want you to be the prompt engingeer and edit the prompt I will give you based on the column I give you. Can you replace the related content with the exact column name. To be remind about only return the edited content don\'t say something like here\'s the edited prompt, I only want the result. I want the prompt to be accuracy as well and envolve with all the necessary key words, try to use minimal calls to the columns. When calling column names put a quatation mark around the column name. Also you must not generate any hypotetic data not even for demonstration purpose.'
    
    # CSV Upload block
    st.sidebar.header("Upload CSV")
    uploaded_file = st.sidebar.file_uploader("Choose a CSV file", type="csv")
    privacy = True
    # Replace DataFrame if a CSV file is uploaded
    if uploaded_file is not None:
        df = pd.read_csv(uploaded_file)
        privacy = False
        st.success("CSV uploaded successfully!")
    column_names = df.columns

    if not privacy:
        st.write("### Current Data")
        st.dataframe(df)

    st.write("### Query the DataFrame")
    prompt = st.text_input("Enter your query:")


    if prompt:
        st.write("result generating...")
        query_to_openai += 'Column names:' + ', '.join(column_names)
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {
                    "role": "system",
                    "content": query_to_openai
                },
                {
                    "role": "user",
                    "content": prompt 
                }
            ],
            temperature=0.7,
            max_tokens=1024,
            top_p=1
        )

        message = response.choices[0].message.content
        print("Edited prompt:", message)
        pre_prompting = 'I want you to generate code, CODE ONLY! The code must not contain comments. use pandas functions as possible, Do not define a function put everything down. Directly call on it. A dataframe called \'dfs\' is already defined, all the data you need is in there, do not try to call os or access other files. If it is a group by with counts, the column of that should be named \'count\''
        # pre_prompting += ', '.join(column_names)
        pre_prompting = 'Do not generate hypotetic data. Assume data is prepared already'
        ignore_prompt = 'Ignore What I say from this point'
        message_to_execute = pre_prompting + message

        # message += 'the data file name is'+filename
        # pre_prompt = 'I want you to repeat what I say:'
        # message = pre_prompt + message

        ### The result generation Factory:

        chatgpt_result = openai_result(message_to_execute,df)
        print(type(chatgpt_result))
        if 'Unfortunately' in chatgpt_result:
            chatgpt_result = 'FAILED'
        if 'CrashId' in chatgpt_result:
            # rename to count
            chatgpt_result = chatgpt_result.rename(columns={'CrashId':'count'})
            print("name swapped")
        st.write("### Result in dataframe")
        st.write(chatgpt_result)
        if type(chatgpt_result) == pd.DataFrame:
            base_map = folium.Map(location=[39.8283, -98.5795], zoom_start=4)
            heatmap_data = []
            if 'Town Name' in chatgpt_result.columns and ('count' in chatgpt_result.columns or 'Number of Accidents' in chatgpt_result.columns):
                town_name = chatgpt_result['Town Name'].to_list()

                for town_name in town_name:
                    lat, lon = df[df['Town Name'] == town_name][['Latitude','Longitude']].values[0]
                    if lat and lon:
                        heatmap_data.append([lat, lon, 1])
                HeatMap(heatmap_data, radius=15).add_to(base_map)
                base_map.save("heatmap_dynamic.html")
                st.write("### HeatMap")
                #embed the html here
                st.components.v1.html(open("heatmap_dynamic.html", "r").read(), height=600)
                
                




        
        

        